import re
from rouge_score import rouge_scorer
import jieba
from nltk.translate.bleu_score import sentence_bleu
from sentence_transformers import SentenceTransformer, util
import numpy as np
from bert_score import score
import torch

def compute_rouge(pred, ref):
    """计算 ROUGE-L 分数"""
    if not pred or not ref:
        return 0.0
    scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=False)
    scores = scorer.score(ref, pred)
    return scores['rougeL'].fmeasure

def compute_bleu(pred, ref):
    """计算 BLEU 分数，使用 jieba 进行中文分词"""
    if not pred or not ref:
        return 0.0
    # 规范化文本，去除多余标点和空格
    def normalize_text(text):
        text = re.sub(r'[，。！？]', '', text)
        text = re.sub(r'\s+', ' ', text)
        return text.strip()
    
    pred = normalize_text(pred)
    ref = normalize_text(ref)
    
    # 使用 jieba 进行分词
    ref_tokens = [list(jieba.cut(ref, cut_all=False))]
    pred_tokens = list(jieba.cut(pred, cut_all=False))
    
    # 计算 BLEU 分数，权重为 4-gram
    try:
        score = sentence_bleu(ref_tokens, pred_tokens, weights=(0.25, 0.25, 0.25, 0.25))
    except ZeroDivisionError:
        score = 0.0
    return score

from bert_score import score

def compute_bert_score(pred: str, ref: str, lang: str = "zh") -> float:
    """
    计算预测文本和参考文本的 BERTScore (F1 分数)。

    Args:
        pred (str): 预测文本
        ref (str): 参考文本
        lang (str): 语言代码，默认为 'zh'（中文）

    Returns:
        float: BERTScore 的 F1 分数
    """
    if not pred or not ref:
        return 0.0
    P, R, F1 = score([pred], [ref], lang=lang, model_type="bert-base-chinese", device="cuda" if torch.cuda.is_available() else "cpu")
    return F1.item()

def compute_safety_completeness(suggestion):
    """检查建议是否包含关键安全措施"""
    if not suggestion:
        return 0.0
    safety_keywords = ["撤离", "通风", "检查", "暂停", "修复", "校准", "更换"]
    score = sum(1 for kw in safety_keywords if kw in suggestion.lower()) / len(safety_keywords)
    return score

def compute_threshold_logic(pred, answer):
    """检查阈值逻辑一致性"""
    pred_is_exceeded = pred.get("is_exceeded", "").strip().lower()
    answer_is_exceeded = answer.get("is_exceeded", "").strip().lower()
    pred_alarm_type = pred.get("alarm_type", "").strip().lower()
    answer_alarm_type = answer.get("alarm_type", "").strip().lower()
    return float(pred_is_exceeded == answer_is_exceeded and pred_alarm_type == answer_alarm_type)

def compute_logic_match(pred, answer):
    """检查整体逻辑匹配"""
    keys = ["is_exceeded", "alarm_type", "alarm_level"]
    matches = sum(1 for key in keys if pred.get(key, "").strip().lower() == answer.get(key, "").strip().lower())
    return matches / len(keys)